"""Encoder registry for stage 2 models."""



# Import modules that perform registration on import.
from typing import Callable, Dict, Optional, Type, Union
from typing import Protocol, Any, runtime_checkable

@runtime_checkable
class Stage2ModelProtocol(Protocol):
    def forward(self, *args: Any, **kwargs: Any) -> Any: ...
    def forward_with_cfg(self, *args: Any, **kwargs: Any) -> Any: ...
    def forward_with_autoguidance(self, *args: Any, **kwargs: Any) -> Any: ...
    def train(self, mode: bool = True) -> Any: ...
    def eval(self) -> Any: ...